#!/usr/bin/env python3
"""
修复chat_sessions表中的novel_id字段
通过character_id从novel_characters表获取对应的novel_id
"""

import asyncio
import sys
import os
from sqlalchemy import text

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from app.database import AsyncSessionLocal


async def check_data():
    """检查当前数据状况"""
    print("=== 检查chat_sessions和novel_characters的关联情况 ===\n")
    
    async with AsyncSessionLocal() as session:
        # 检查chat_sessions的character_id和novel_id情况
        result = await session.execute(text("""
            SELECT 
                cs.id as session_id,
                cs.character_id,
                cs.novel_id,
                nc.novel_id as character_novel_id,
                nc.name as character_name
            FROM chat_sessions cs
            LEFT JOIN novel_characters nc ON cs.character_id = nc.id
            ORDER BY cs.id
        """))
        rows = result.fetchall()
        
        print("当前chat_sessions和novel_characters的关联情况:")
        for row in rows:
            session_id, char_id, session_novel_id, char_novel_id, char_name = row
            print(f"  Session ID: {session_id}")
            print(f"    Character ID: {char_id}")
            print(f"    Session Novel ID: {session_novel_id}")
            print(f"    Character Novel ID: {char_novel_id}")
            print(f"    Character Name: {char_name}")
            print()
        
        # 检查有多少个sessions需要修复
        result = await session.execute(text("""
            SELECT COUNT(*) as count
            FROM chat_sessions cs
            JOIN novel_characters nc ON cs.character_id = nc.id
            WHERE cs.novel_id IS NULL AND nc.novel_id IS NOT NULL
        """))
        count = result.fetchone()[0]
        print(f"需要修复的chat_sessions记录数: {count}")
        
        return count


async def fix_novel_id():
    """修复novel_id字段"""
    print("\n=== 开始修复novel_id字段 ===\n")
    
    async with AsyncSessionLocal() as session:
        # 执行修复
        result = await session.execute(text("""
            UPDATE chat_sessions cs
            JOIN novel_characters nc ON cs.character_id = nc.id
            SET cs.novel_id = nc.novel_id
            WHERE cs.novel_id IS NULL AND nc.novel_id IS NOT NULL
        """))
        
        affected_rows = result.rowcount
        await session.commit()
        
        print(f"修复完成，更新了 {affected_rows} 条记录")
        
        return affected_rows


async def verify_fix():
    """验证修复结果"""
    print("\n=== 验证修复结果 ===\n")
    
    async with AsyncSessionLocal() as session:
        # 检查修复后的数据
        result = await session.execute(text("""
            SELECT 
                cs.id as session_id,
                cs.character_id,
                cs.novel_id,
                nc.novel_id as character_novel_id,
                nc.name as character_name
            FROM chat_sessions cs
            LEFT JOIN novel_characters nc ON cs.character_id = nc.id
            ORDER BY cs.id
        """))
        rows = result.fetchall()
        
        print("修复后的数据状况:")
        null_count = 0
        for row in rows:
            session_id, char_id, session_novel_id, char_novel_id, char_name = row
            print(f"  Session ID: {session_id}, Novel ID: {session_novel_id}, Character: {char_name}")
            if session_novel_id is None:
                null_count += 1
        
        print(f"\n还有 {null_count} 条记录的novel_id为NULL")
        
        # 检查是否还有不匹配的情况
        result = await session.execute(text("""
            SELECT COUNT(*) as count
            FROM chat_sessions cs
            JOIN novel_characters nc ON cs.character_id = nc.id
            WHERE cs.novel_id != nc.novel_id
        """))
        mismatch_count = result.fetchone()[0]
        print(f"有 {mismatch_count} 条记录的novel_id与character的novel_id不匹配")


async def main():
    try:
        # 1. 检查数据状况
        count_to_fix = await check_data()
        
        if count_to_fix == 0:
            print("所有记录的novel_id都已正确设置，无需修复。")
            return
        
        # 2. 询问用户是否继续修复
        response = input(f"\n发现 {count_to_fix} 条记录需要修复，是否继续？(y/N): ")
        if response.lower() != 'y':
            print("修复已取消。")
            return
        
        # 3. 执行修复
        fixed_count = await fix_novel_id()
        
        # 4. 验证修复结果
        await verify_fix()
        
        print(f"\n✅ 修复完成！成功更新了 {fixed_count} 条记录。")
        
    except Exception as e:
        print(f"❌ 修复过程中出现错误: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(main()) 